import chisel3._
import chisel3.util._
import Control._

class IFU(xlen:Int) extends Module {
  val io = IO(new Bundle {
    val done = Output(Bool())
    val addr = Output(UInt(xlen.W))

    val in  = Flipped(Decoupled(new MessageWI(xlen)))
    val out = Decoupled(new MessageFD(xlen))

    // IFU <=> ICache
    val data  = Input(UInt(xlen.W))
    val avalid= Output(Bool())
    val aready= Input(Bool())
    val rvalid= Input(Bool())
    val rready= Output(Bool())

    // val valid = Output(Bool()) // avalid
    // val ready = Input(Bool())  // rvalid

    // IFU <=> Master interface
    val ACLK    = Input(Clock())
    val ARESETn = Input(Reset())

    // IFU <=> MEU
    val br_taken = Input(Bool())
    val Alu_out  = Input(UInt(xlen.W))
  })

  val pc = Module(new PC(xlen, xlen))

  val inst       = Wire(UInt(xlen.W))
  val inst_valid = Wire(Bool())
  val ren        = Wire(Bool())
  val done       = Wire(Bool())
  val areset_pos = Wire(Bool())
  // val flush      = RegInit(0.B)
  // val reading    = RegInit(0.B)
  val rready     = RegInit(1.B)

  // flush := Mux(reading & io.rvalid & io.rready, 0.B, Mux(io.br_taken, 1.B, flush))
  // reading := Mux(io.rvalid & io.rready, 0.B, Mux(io.avalid & io.aready, 1.B, reading))

  /* state machine */
  val sIDEL :: sIREAD :: sFLUSH :: sREREAD :: Nil = Enum(4)
  val curr_status = RegInit(sIDEL)
  val next_status = WireDefault(sIDEL)

  // status transfer
  curr_status := next_status
  switch (curr_status) {
    is (sIDEL) {
      when(areset_pos || (io.out.valid & io.out.ready) || io.br_taken) {
        next_status := sIREAD;
      } .otherwise {
        next_status := sIDEL;
      }
    }

    // read instruction
    is (sIREAD) {
      when (io.br_taken & io.rready & io.rvalid) {
        next_status := sREREAD;
      } .elsewhen (io.br_taken) {
        next_status := sFLUSH;
      } .elsewhen(io.rready & io.rvalid) {
        next_status := sIDEL;
      } .otherwise {
        next_status := sIREAD;
      }
    }

    // flush
    is (sFLUSH) {
      when (io.rready & io.rvalid) {
        next_status := sREREAD;
      } .otherwise {
        next_status := sFLUSH;
      }
    }

    // read again
    is (sREREAD) {
      when (io.rready & io.rvalid) {
        next_status := sIDEL;
      } .otherwise {
        next_status := sREREAD;
      }
    }
  }

  // synchronize with outside clock
  withClockAndReset(io.ACLK, !(io.ARESETn).asBool) {
    val _inst = RegInit(0.U(xlen.W))
    val _inst_valid = RegInit(0.B)

    _inst := Mux(io.rready & io.rvalid, io.data, _inst)
    // _inst_valid := Mux(!_inst_valid && io.rready, Mux(flush | io.br_taken, 0.B, 1.B), Mux(io.out.ready && io.out.valid, 0.B, _inst_valid))
    // _inst_valid := Mux(_inst_valid & (flush | io.br_taken), 0.B, Mux(io.rready & io.rvalid, Mux(flush | io.br_taken, 0.B, 1.B), Mux(io.out.ready && io.out.valid, 0.B, _inst_valid))) // BUG: 情况似乎没有考虑完整
    _inst_valid := Mux(io.out.ready & io.out.valid | io.br_taken, 0.B, Mux(curr_status===sREREAD && io.rready && io.rvalid || curr_status===sIREAD && io.rready && io.rvalid && ~io.br_taken, 1.B, _inst_valid))

    val ARESETn_raise = Wire(Bool())
    val ARESETn_cur   = RegInit(0.B)
    val ARESETn_pre   = RegInit(0.B)

    ARESETn_cur   := io.ARESETn.asBool
    ARESETn_pre   := ARESETn_cur
    ARESETn_raise := ~ARESETn_pre & ARESETn_cur

    val _ren = RegInit(0.B)
    // _ren := Mux(ARESETn_raise, 1.B, Mux(io.in.ready && io.in.valid, 1.B, 0.B))
    // _ren := Mux(ARESETn_raise, 1.B, Mux(io.ready, 1.B, 0.B))
    // _ren := Mux(ARESETn_raise, 1.B, Mux((flush | io.br_taken) & reading & io.rready, 1.B, Mux(io.out.ready & io.out.valid, 1.B, 0.B)))
    // _ren := Mux(io.aready & io.avalid, 0.B, Mux(ARESETn_raise, 1.B, Mux(io.out.ready & io.out.valid, 1.B, _ren)))
    // _ren := Mux(io.aready & io.avalid, 0.B, Mux(flush | io.br_taken, Mux(reading, Mux(io.rready & io.rvalid, 1.B, _ren), 1.B), Mux(ARESETn_raise, 1.B, Mux(io.out.ready & io.out.valid, 1.B, _ren))))
    // _ren := Mux(io.aready & io.avalid, 0.B, Mux(ARESETn_raise, 1.B, Mux((flush | io.br_taken) & reading & io.rready & io.rvalid, 1.B, Mux(io.out.ready & io.out.valid, 1.B, 0.B))))
    _ren := Mux(curr_status===sIDEL && next_status===sIREAD || curr_status===sFLUSH && next_status===sREREAD || curr_status===sIREAD && next_status===sREREAD, 1.B, 0.B)

    val _done = RegInit(0.B)
    _done := Mux(_inst_valid && io.in.valid, 1.B, 0.B)

    inst       := _inst
    inst_valid := _inst_valid
    ren        := _ren
    done       := _done
    areset_pos := ARESETn_raise
  }


  io.addr              := Mux(io.br_taken, io.Alu_out, pc.io.pc_out)

  io.out.bits.pc       := pc.io.pc_out
  io.out.bits.inst     := inst

  // pc.io.PC_sel   := Mux(io.br_taken, PC_ALU, io.in.bits.PC_sel)
  pc.io.PC_sel   := Mux(io.br_taken, PC_ALU, PC_4)
  pc.io.br_taken := io.br_taken
  pc.io.csr_out  := io.in.bits.csr_out
  pc.io.Alu_out  := io.Alu_out
  // pc.io.wen      := io.in.ready & io.in.valid
  pc.io.wen      := (io.out.ready & io.out.valid) | io.br_taken // pc + 4 for every instruction fetched

  io.done := done

  // io.valid := ren
  io.avalid := ren

  // rready := Mux(io.rready & io.rvalid, Mux(flush | io.br_taken, 1.B, 0.B), Mux(io.out.ready & io.out.valid, 1.B, rready))
  rready := Mux(io.rready & io.rvalid, 0.B, Mux(curr_status===sIREAD || curr_status===sREREAD, 1.B, rready))
  io.rready := rready

  io.out.valid := inst_valid
  io.in.ready  := 1.B
}
